import os
import sys
import shutil
import subprocess
import tempfile
from pathlib import Path
from loguru import logger

def compile_py_to_so(
    py_file: str,
    output_dir: str = None,
    optimize_level: int = 2,
    python_version: str = "3",
    keep_source: bool = False,
    delete_temp: bool = True,
    extra_args: str = None,
    log_callback = None
):
    """
    将Python文件编译为.so二进制模块
    
    参数:
        py_file: Python文件路径
        output_dir: 输出目录，默认为源文件所在目录
        optimize_level: 优化级别 0-3
        python_version: Python版本 "2" 或 "3"
        keep_source: 是否保留源文件
        delete_temp: 是否删除临时文件
        extra_args: 额外的编译参数
        log_callback: 日志回调函数，用于更新进度
    
    返回:
        生成的.so文件路径
    """
    def log(message):
        """日志辅助函数"""
        if log_callback:
            log_callback(message)
        logger.info(message)
    
    py_file = Path(py_file).resolve()
    if not py_file.exists():
        error_msg = f"文件不存在: {py_file}"
        log(f"[ERROR] {error_msg}")
        raise Exception(error_msg)
    
    if py_file.suffix != ".py":
        error_msg = f"仅支持 .py 文件，不支持 {py_file.suffix}"
        log(f"[ERROR] {error_msg}")
        raise Exception(error_msg)
    
    # 设置输出目录
    if output_dir:
        output_dir = Path(output_dir).resolve()
        output_dir.mkdir(parents=True, exist_ok=True)
    else:
        output_dir = py_file.parent
    
    module_name = py_file.stem
    log(f"开始编译模块: {module_name}")
    
    # 创建临时目录用于编译
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_dir = Path(tmp_dir)
        tmp_py = tmp_dir / f"{module_name}.py"
        tmp_setup = tmp_dir / "setup.py"
        
        # 复制源文件到临时目录
        shutil.copy(py_file, tmp_py)
        log(f"复制源文件到临时目录: {tmp_py}")
        
        # 生成 setup.py 内容
        setup_code = f'''
from setuptools import setup
from Cython.Build import cythonize

setup(
    ext_modules=cythonize(
        "{module_name}.py",
        compiler_directives={{"language_level": {python_version}}},
        build_dir="build"
    )
)
'''
        tmp_setup.write_text(setup_code)
        log("生成 setup.py 文件")
        
        # 构建编译命令
        cmd = [sys.executable, str(tmp_setup), "build_ext", "--inplace"]
        if extra_args:
            cmd.extend(extra_args.split())
        
        log(f"执行编译命令: {' '.join(cmd)}")
        result = subprocess.run(
            cmd,
            cwd=tmp_dir,
            capture_output=True,
            text=True
        )
        
        # 检查编译结果
        if result.returncode != 0:
            error_msg = f"编译失败: {result.stderr}"
            log(f"[ERROR] {error_msg}")
            raise Exception(error_msg)
        
        # 查找生成的 .so 文件
        so_files = list(tmp_dir.glob(f"{module_name}*.so"))
        if sys.platform.startswith('win'):
            # Windows 平台上是 .pyd 文件
            so_files = list(tmp_dir.glob(f"{module_name}*.pyd"))
        
        if not so_files:
            error_msg = "未找到生成的二进制文件"
            log(f"[ERROR] {error_msg}")
            raise Exception(error_msg)
        
        so_file = so_files[0]
        log(f"找到生成的二进制文件: {so_file.name}")
        
        # 移动到输出目录
        target_so = output_dir / so_file.name
        shutil.move(str(so_file), target_so)
        log(f"已将二进制文件移动到输出目录: {target_so}")
    
    # 清理中间文件
    c_file = output_dir / f"{module_name}.c"
    if c_file.exists():
        c_file.unlink()
        log(f"已删除中间文件: {c_file.name}")
    
    # 清理 build 目录
    build_dir = output_dir / "build"
    if build_dir.exists() and build_dir.is_dir():
        shutil.rmtree(build_dir)
        log(f"已删除构建目录: {build_dir.name}")
    
    # 删除源文件（如果需要）
    if not keep_source and py_file.exists():
        py_file.unlink()
        log(f"已删除源文件: {py_file.name}")
    elif keep_source:
        log(f"保留源文件: {py_file.name}")
    
    log(f"编译完成: {target_so}")
    return str(target_so)
